#!/usr/bin/env python3
import os
import re
import sys
from cpp_flow_map.core import analyze_cpp_code

def fix_main_and_run(src_dir, target_file):
    """
    A wrapper that finds the main function first, then runs the analysis.
    """
    # First, find the main function path
    # Convert src_dir to a list if it's not already
    src_directories = src_dir if isinstance(src_dir, list) else [src_dir]
    
    # Process each directory
    main_file = None
    for directory in src_directories:
        main_file = find_main_file(directory)
        if main_file:
            break
    
    if main_file:
        # Now force this as the main entry point before analysis
        print(f"Found main function in {main_file}")
        os.environ['CPP_FLOW_MAP_MAIN_FILE'] = main_file
        
        # Patch the core module's entry point detection
        from cpp_flow_map.Tools.CppCodeParser import CppCodeParser
        original_find_entry = CppCodeParser._find_entry_points
        
        def patched_find_entry(self):
            # Call original method
            original_find_entry(self)
            
            # Check if we have a known main file
            main_file_path = os.environ.get('CPP_FLOW_MAP_MAIN_FILE')
            if main_file_path:
                main_base = os.path.splitext(os.path.basename(main_file_path))[0]
                
                # If we don't have a main entry point yet, check for the main file
                main_found = False
                for entry in self.entry_points:
                    if entry == "main":
                        main_found = True
                        break
                
                if not main_found:
                    # Try to add main function to entry points
                    for func_name, func_node in self.functions.items():
                        if func_name == "main":
                            if func_name not in self.entry_points:
                                self.entry_points.insert(0, func_name)
                                print(f"Added main function as entry point")
                                main_found = True
                                break
                
                # Reorder entry points to prioritize main
                if "main" in self.entry_points:
                    self.entry_points.remove("main")
                    self.entry_points.insert(0, "main")
                    main_found = True
                
                if not main_found:
                    print(f"Warning: Despite finding main.c, couldn't add main function as entry point")
        
        # Apply the patch
        CppCodeParser._find_entry_points = patched_find_entry
        
        # Run the analysis
        print(f"Running analysis with {src_directories} and {target_file}...")
        analyze_cpp_code(src_directories, target_file)
    else:
        print("Could not find main function. Running normal analysis...")
        analyze_cpp_code(src_dir, target_file)

def find_main_file(directory):
    """Find the file containing the main function."""
    # Main function patterns
    main_patterns = [
        r'int\s+main\s*\(\s*int\s+\w+\s*,\s*char\s*\*+\s*\w+\[\s*\]\s*\)\s*\{',  # int main(int argc, char *argv[]) {
        r'int\s+main\s*\(\s*int\s+\w+\s*,\s*char\s*\*+\s*\*+\s*\w+\s*\)\s*\{',    # int main(int argc, char **argv) {
        r'int\s+main\s*\(\s*void\s*\)\s*\{',                                     # int main(void) {
        r'int\s+main\s*\(\s*\)\s*\{',                                            # int main() {
        r'int\s+main\s*\(\s*int\s+\w+\s*\)\s*\{'                                 # int main(int argc) {
    ]
    
    # Find all C/C++ files
    cpp_files = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith(('.c', '.cpp', '.cxx', '.cc')):
                cpp_files.append(os.path.join(root, file))
    
    print(f"Searching for main function in {len(cpp_files)} files...")
    
    # First, check files named main.c or main.cpp as they're most likely to contain main
    main_files = [f for f in cpp_files if os.path.basename(f).lower().startswith('main.')]
    if main_files:
        for file_path in main_files:
            if has_main_function(file_path, main_patterns):
                return file_path
    
    # If not found in main.* files, check all other files
    for file_path in cpp_files:
        if has_main_function(file_path, main_patterns):
            return file_path
    
    return None

def has_main_function(file_path, patterns):
    """Check if a file contains a main function."""
    try:
        with open(file_path, 'r', encoding='utf-8', errors='ignore') as file:
            content = file.read()
            
            for pattern in patterns:
                if re.search(pattern, content, re.MULTILINE):
                    print(f"Found main function in {file_path}")
                    return True
    except Exception as e:
        print(f"Error reading file {file_path}: {e}")
    
    return False

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python fix_main_detection.py <source_directory> [target_file]")
        sys.exit(1)
    
    src_dir = sys.argv[1]  # Pass the directory as a string
    target_file = sys.argv[2] if len(sys.argv) > 2 else "main.c"
    
    fix_main_and_run(src_dir, target_file) 